package com.cat.dataStructure02;

import java.util.HashMap;
import java.util.Map;

/**
 * @author 曲大人的喵
 * @description https://leetcode.cn/problems/count-subarrays-with-median-k/description/
 * @create 2025/10/18 21:16
 * @since JDK17
 */

public class Solution21 {
    public int countSubarrays(int[] nums, int k) {
        int ans = 0, n = nums.length;
        Map<Integer, Integer> map = new HashMap<>();
        int pos = 0;
        while (nums[pos] != k) {
            pos++;
        }
        map.put(0, 1);
        for (int i = pos - 1, cnt = 0; i >= 0; i--) {
            cnt += nums[i] > k ? -1 : 1;
            map.merge(cnt, 1, Integer::sum);
        }
        ans += map.get(0) + map.getOrDefault(-1, 0);
        for (int i = pos + 1, cnt = 0; i < n; i++) {
            cnt += nums[i] > k ? 1 : -1;
            ans += map.getOrDefault(cnt - 1, 0) + map.getOrDefault(cnt, 0);
        }

        return ans;
    }
}
